#include "mlp_dcu.h"

// 矩阵乘法核函数 - 优化版本
// 使用分块策略以提高缓存命中率
__global__ void matmul_kernel(const double* A, const double* B, double* C, int M, int N, int K) {
    // 定义共享内存块大小
    constexpr int BLOCK_SIZE = 16;
    
    // 声明共享内存
    __shared__ double As[BLOCK_SIZE][BLOCK_SIZE];
    __shared__ double Bs[BLOCK_SIZE][BLOCK_SIZE];
    
    // 计算线程和块索引
    int bx = blockIdx.x;
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    
    // 计算全局行列索引
    int row = by * BLOCK_SIZE + ty;
    int col = bx * BLOCK_SIZE + tx;
    
    // 累加器
    double sum = 0.0;
    
    // 遍历所有子块
    for (int i = 0; i < (K + BLOCK_SIZE - 1) / BLOCK_SIZE; ++i) {
        // 协作加载A和B的子块到共享内存
        if (row < M && i * BLOCK_SIZE + tx < K) {
            As[ty][tx] = A[row * K + i * BLOCK_SIZE + tx];
        } else {
            As[ty][tx] = 0.0;
        }
        
        if (i * BLOCK_SIZE + ty < K && col < N) {
            Bs[ty][tx] = B[(i * BLOCK_SIZE + ty) * N + col];
        } else {
            Bs[ty][tx] = 0.0;
        }
        
        // 同步以确保数据加载完成
        __syncthreads();
        
        // 计算当前子块的点积
        for (int k = 0; k < BLOCK_SIZE; ++k) {
            sum += As[ty][k] * Bs[k][tx];
        }
        
        // 同步以确保计算完成
        __syncthreads();
    }
    
    // 写入结果
    if (row < M && col < N) {
        C[row * N + col] = sum;
    }
}

// 添加偏置核函数
__global__ void add_bias_kernel(double* output, const double* bias, int rows, int cols) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (row < rows && col < cols) {
        output[row * cols + col] += bias[col];
    }
}

// ReLU激活函数核函数
__global__ void relu_kernel(double* data, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        data[idx] = fmax(0.0, data[idx]);
    }
}

// LinearLayer的前向传播实现
void LinearLayer::forward(const Matrix& input, Matrix& output) {
    int batch_size = input.getRows();
    
    // 确保输出矩阵大小正确
    if (output.getRows() != batch_size || output.getCols() != output_dim) {
        // 如果输出矩阵尺寸不匹配，重新创建
        output = Matrix(batch_size, output_dim);
    }
    
    // 分配设备内存
    output.allocateDevice();
    
    // 执行矩阵乘法
    matrixMultiply(input, weights, output);
    
    // 添加偏置
    addBias(output, bias);
}

// 矩阵乘法实现
void LinearLayer::matrixMultiply(const Matrix& A, const Matrix& B, Matrix& C) {
    int M = A.getRows();    // 输入批次大小
    int K = A.getCols();    // 输入特征维度
    int N = B.getCols();    // 输出特征维度
    
    // 定义块和线程大小
    constexpr int BLOCK_SIZE = 16;
    dim3 threadsPerBlock(BLOCK_SIZE, BLOCK_SIZE);
    dim3 blocksPerGrid((N + BLOCK_SIZE - 1) / BLOCK_SIZE, 
                       (M + BLOCK_SIZE - 1) / BLOCK_SIZE);
    
    // 启动核函数
    matmul_kernel<<<blocksPerGrid, threadsPerBlock>>>(
        A.getDeviceData(), B.getDeviceData(), C.getDeviceData(), M, N, K);
    
    // 检查错误
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipDeviceSynchronize());
}

// 添加偏置实现
void LinearLayer::addBias(Matrix& output, const Matrix& bias) {
    int rows = output.getRows();
    int cols = output.getCols();
    
    // 定义块和线程大小
    constexpr int BLOCK_SIZE = 16;
    dim3 threadsPerBlock(BLOCK_SIZE, BLOCK_SIZE);
    dim3 blocksPerGrid((cols + BLOCK_SIZE - 1) / BLOCK_SIZE, 
                       (rows + BLOCK_SIZE - 1) / BLOCK_SIZE);
    
    // 启动核函数
    add_bias_kernel<<<blocksPerGrid, threadsPerBlock>>>(
        output.getDeviceData(), bias.getDeviceData(), rows, cols);
    
    // 检查错误
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipDeviceSynchronize());
}

// ReLULayer的前向传播实现
void ReLULayer::forward(const Matrix& input, Matrix& output) {
    // 如果输入和输出不是同一个矩阵，复制输入到输出
    if (&input != &output) {
        // 确保输出矩阵大小正确
        if (output.getRows() != input.getRows() || output.getCols() != input.getCols()) {
            // 如果输出矩阵尺寸不匹配，重新创建
            output = Matrix(input.getRows(), input.getCols());
        }
        
        // 分配设备内存
        output.allocateDevice();
        
        // 复制输入数据到输出
        HIP_CHECK(hipMemcpy(
            output.getDeviceData(), 
            input.getDeviceData(), 
            input.size() * sizeof(double), 
            hipMemcpyDeviceToDevice));
    }
    
    // 应用ReLU激活函数
    applyReLU(output);
}

// 应用ReLU激活函数
void ReLULayer::applyReLU(Matrix& data) {
    int size = data.size();
    
    // 定义块和线程大小
    constexpr int BLOCK_SIZE = 256;
    int blocks = (size + BLOCK_SIZE - 1) / BLOCK_SIZE;
    
    // 启动核函数
    relu_kernel<<<blocks, BLOCK_SIZE>>>(data.getDeviceData(), size);
    
    // 检查错误
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipDeviceSynchronize());
}
